'''
https://leetcode.cn/problems/smallest-range-covering-elements-from-k-lists/description/
'''
import heapq
from typing import List
from sortedcontainers import SortedDict, SortedSet

class Solution:
    # 有序表
    def smallestRange(self, nums: List[List[int]]) -> List[int]:
        m = len(nums)
        # element: (value, i, j)  where i: nums[i], j:nums[i][j], v=nums[i][j]
        ss = SortedSet([(nums[i][0], i, 0) for i, lists in enumerate(nums)])

        min_v, max_v = ss[0][0], ss[-1][0]
        while len(ss) == m:
            max_v_in_set = ss[-1][0]
            min_v_in_set, i, j = ss.pop(0)
            if max_v_in_set - min_v_in_set < max_v - min_v:
                min_v, max_v = min_v_in_set, max_v_in_set
            if j + 1 < len(nums[i]):
                ss.add( (nums[i][j+1], i, j + 1) )
        return [min_v, max_v]

    # 使用堆实现，逻辑一样
    def smallestRange2(self, nums: List[List[int]]) -> List[int]:
        m = len(nums)
        # element: (value, i, j)  where i: nums[i], j:nums[i][j], v=nums[i][j]
        heap = [(nums[i][0], i, 0) for i, lists in enumerate(nums)]
        heapq.heapify(heap)

        min_v_in_heap, max_v_in_heap  = heap[0][0], max(heap)[0]
        min_v, max_v = min_v_in_heap, max_v_in_heap
        while len(heap) == m:
            min_v_in_heap, i, j = heapq.heappop(heap)
            if max_v_in_heap - min_v_in_heap < max_v - min_v:
                min_v, max_v = min_v_in_heap, max_v_in_heap
            if j + 1 < len(nums[i]):
                heapq.heappush(heap, (nums[i][j+1], i, j + 1))
                max_v_in_heap = max(max_v_in_heap, nums[i][j+1])
        return [min_v, max_v]



nums = [[4,10,15,24,26], [0,9,12,20], [5,18,22,30]]
res = Solution().smallestRange(nums)
print(res)